from langchain_core.prompts import ChatPromptTemplate,FewShotChatMessagePromptTemplate,SystemMessagePromptTemplate
from langchain_ollama import OllamaLLM
from langchain_core.output_parsers import StrOutputParser

llm = OllamaLLM(
    model="deepseek-r1:14b",  # 与 ollama list 中的模型名一致
    base_url="http://localhost:11434",  # Ollama 服务地址
    temperature=0.7,          # 控制随机性 (0-1)
    num_predict=1024,         # 最大生成长度
    top_k=50,                 # 采样参数
    top_p=0.9,                # 采样参数
    stop=["<|EOT|>"],         # DeepSeek 的停止标记
    num_gpu=99,               # 使用全部 GPU 层加速
    num_thread=8,             # CPU 线程数 (4090 可设 8-12)
)

# 1. 定义少样本示例（小红书标题+话题案例）
examples = [
    {
        "theme": "Python学习", 
        "title": "7天搞定Python数据分析！小白逆袭指南🔥",
        "topics": "#Python入门 #数据分析 #编程学习"
    },
    {
        "theme": "考研英语", 
        "title": "考研英语90分秘籍：阅读理解满分技巧💯",
        "topics": "#考研英语 #阅读理解 #学习干货"
    },
    {
        "theme": "高效笔记", 
        "title": "iPad无纸化学习｜Notion模板免费分享✨",
        "topics": "#无纸化学习 #学生党必备 #学习工具"
    }
]

# 2. 定义单个示例的格式化模板
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "主题：{theme}"),
    ("ai", "标题：{title}\n话题：{topics}")
])

# 3. 创建少样本消息模板
few_shot_template = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
    input_variables=["theme"]
)

# 4. 构建完整聊天提示（整合系统角色 + 少样本示例 + 用户输入）
final_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(
        "你是一个资深小红书学习博主，擅长用吸引眼球的标题和话题推广知识。"
        "根据主题生成符合年轻人喜好的标题（带emoji🔥✨💯），并添加3个相关话题（以#开头）"
    ),
    few_shot_template,  # 动态插入少样本示例
    ("human", "主题：{theme}")  # 用户指定本次生成的主题
])

# 5. 创建链
chain = final_prompt | llm | StrOutputParser()

# 6. 执行调用
result = chain.invoke({"theme": "雅思备考"})
print(result)